#include "image_extractor.h"
#include "common.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

const char PNG_SOF[4] = { '\x89', 'P', 'N', 'G' };
const char PNG_EOF[8] = { 'I', 'E', 'N', 'D', '\xAE', 'B', '`', '\x82' };

long get_start_position(FILE *src, unsigned int sn) {
	// save the current position indicator.
	long current = ftell(src);
	// move the file position indicator to the beginning of src.
	rewind(src);

	if (src == NULL) {
		return -1;
	}
	char buffer[4];
	int count = 0;
	while (fread(buffer, 1, 4, src) == 4) {
		// if matched, check count equals sn or not.
		if (is_png_start(buffer) == 0) {
			++count;
			if (count == sn) {
				// move file position indicator to the position saved above.
				long value = ftell(src) - 4;
				fseek(src, current, SEEK_SET);
				return value;
			}
		} else {
			// fallback 3 byte otherwise, you may miss one image-head tag,
			// because the rest bytes of the block you readed may match with
			// the beginning part of the image-head tag.
			fseek(src, -3, SEEK_CUR);
		}
	}

	return -1;
}

long get_end_position(FILE *src, unsigned int sn) {
	if (src == NULL) {
		return -1;
	}
	// save the current position indicator.
	long current = ftell(src);
	// move the file position indicator to the beginning of src.
	rewind(src);

	char buffer[8];
	int count = 0;
	while (fread(buffer, 1, 8, src) == 8) {
		// if matched, check count equals sn or not.
		if (is_png_end(buffer) == 0) {
			++count;
			if (count == sn) {
				long value = ftell(src);
				fseek(src, current, SEEK_SET);
				return value - 1;
			}
		} else {
			// fallback 7 byte otherwise, you may miss one image-tail tag,
			// because the rest bytes of the block you readed may match with
			// the beginning part of the image-tail tag.
			fseek(src, -7, SEEK_CUR);
		}
	}

	return -1;
}

void get_position(FILE *src, unsigned int sn, long *start_pos, long *end_pos) {
}

int is_png_start(const char *block) {
	// strcmp can not compare two char array without '\0'
	// so use memcmp.
	return memcmp(block, PNG_SOF, 4);
}

int is_png_end(const char *block) {
	return memcmp(block, PNG_EOF, 8);
}

int extract_all_file(FILE *src, const char* out_dir) {
	// malloc memory for file names;
	int count = get_picture_count(src);
	char **file_names = (char**) malloc(sizeof(char*) * count);

	// read the end position of resources.
	int res_end_pos;
	fseek(src, 16, SEEK_SET);
	fread((void*) &res_end_pos, 4, 1, src);
	res_end_pos = reverse_int32(res_end_pos);
	// move to the next byte of that position
	fseek(src, res_end_pos, SEEK_SET);

	// start to read file names
	int i;
	for (i = 0; i < count; ++i) {
		// read the length of the current file name.
		short name_length = 0;
		fread((void*) &name_length, 2, 1, src);
		name_length = reverse_int16(name_length);
		file_names[i] = (char *) malloc(sizeof(char) * (name_length + 1));
		// skip 4 unknown bytes
		fseek(src, 4, SEEK_CUR);
		// read the file name.
		int j;
		for (j = 0; j < name_length; ++j) {
			fseek(src, 1, SEEK_CUR);
			fread((void*) (&file_names[i][j]), sizeof(char), 1, src);
		}
		file_names[i][name_length] = '\0';
	}
	rewind(src);
	for (i = 0; i < count; ++i) {
		// generate the full path of output file
		char file_path[255] = "\0";
		char sn_buffer[16];
		strcat(file_path, out_dir);
		strcat(file_path, itoa(i, sn_buffer, 10));
		strcat(file_path, "-");
		strcat(file_path, file_names[i]);
		FILE *out = fopen(file_path, "wb");
		rewind(src);
		if (extract_file(src, i + 1, out) != 0) {
			remove(file_path);
		}
		free(file_names[i]);
	}
	free(file_names);
	return 0;
}
int extract_file(FILE *src, unsigned int sn, FILE *out) {
	// save the current position indicator.
	long current = ftell(src);
	// move the file position indicator to the beginning of src.
	rewind(src);

	long start_pos = get_start_position(src, sn);
	long end_pos = get_end_position(src, sn);

	int length = end_pos - start_pos + 1;
	if (length < 0) {
		fseek(src, current, SEEK_SET);
		return -1;
	}
	fseek(src, start_pos, SEEK_SET);

	char *buffer = (char*) malloc(length);
	if (fread(buffer, 1, length, src) == length) {
		fwrite(buffer, 1, length, out);
	} else {
		fseek(src, current, SEEK_SET);
		free(buffer);
		return -1;
	}

	fseek(src, current, SEEK_SET);
	free(buffer);
	return 0;
}

int get_picture_count(FILE *src) {
	if (src == NULL) {
		return -1;
	}
	// save the current position indicator.
	long current = ftell(src);
	// move the file position indicator to the beginning of src.
	rewind(src);

	char buffer[4];
	int count = 0;
	while (fread(buffer, 1, 4, src) == 4) {
		if (is_png_start(buffer) == 0) {
			++count;
		} else {
			fseek(src, -3, SEEK_CUR);
		}
	}

	fseek(src, current, SEEK_SET);
	return count;
}

int get_valid_length(FILE *src) {
	fseek(src, 8, SEEK_SET);
	int value;
	fread(&value, 4, 1, src);
	return reverse_int32(value);
}

int replace(FILE *src, int sn, FILE *new_picture, FILE *out) {
	if (src == NULL, new_picture == NULL, out == NULL) {
		return -1;
	}
	// mark the previous position indicator.
	long previous_pos = ftell(src);

	// get the old picture size.
	int old_size = get_picture_size(src, sn);
	// get the new picture's size
	fseek(new_picture, 0, SEEK_END);
	int new_picture_size = ftell(new_picture);
	// calculate the change of the Valid Data Block(VDB) size
	int size_changed = new_picture_size - old_size;
	// replace the head info.
	if (size_changed != 0) {
		// copy the first 8 bytes to out file
		if (copy_part(src, 0, 7, out) == -1) {
			return -1;
		}
		// read the old VDB size
		fseek(src, 8, SEEK_SET);
		int old_vdb_size;
		fread((void*) &old_vdb_size, 4, 1, src);
		old_vdb_size = reverse_int32(old_vdb_size);
		// calculate the new VDB size
		int new_vdb_size = old_vdb_size + size_changed;
		// write the new VDB size to out file.
		write_int(out, new_vdb_size, 0);
		// copy the next 4 bytes to out file(head size).
		copy_part(src, 12, 15, out);
		// read the old head+pictures_sequence(HPS) size.
		fseek(src, 16, SEEK_SET);
		int old_hps_size;
		fread((void*) &old_hps_size, 4, 1, src);
		old_hps_size = reverse_int32(old_hps_size);
		// calculate the new HPS size.
		int new_hps_size = old_hps_size + size_changed;
		// write the new HPS size to out file.
		write_int(out, new_hps_size, 0);
	} else {
		// copy the first 20 bytes(all the header information) to output file.
		if (copy_part(src, 0, 19, out) == -1) {
			return -1;
		}
	}
	long start_pos = get_start_position(src, sn);
	long end_pos = get_end_position(src, sn);
	if (start_pos > end_pos) {
		return -1;
	}

	if (size_changed != 0) {
		int length = end_pos - start_pos + 1;

		// copy the bytes between file head and the sn-th picture to out.
		// because there are 4 bytes before the picture which mark the picture's size,
		// the final start position is start_pos - 5.
		if (copy_part(src, 20, start_pos - 5, out) == -1) {
			return -1;
		}

		// write the new_picture_size to the output file.
		write_int(out, new_picture_size, 0);
	} else {
		// copy the content between the head and the old picture to output file.
		if (copy_part(src, 20, start_pos - 1, out) == -1) {
			return -1;
		}
	}
	// copy the bytes of the new picture to out.
	if (copy_part(new_picture, 0, new_picture_size - 1, out) == -1) {
		return -1;
	}

	if (size_changed != 0) {
		// write the rest of the valid block to out.
		int valid_length = get_valid_length(src);
		if (copy_part(src, end_pos + 1, valid_length - 1, out) == -1) {
			return -1;
		}

		// write the first 14 bytes
		if (copy_part(src, valid_length, valid_length + 13, out) == -1) {
			return -1;
		}

		// read index array
		int count = get_picture_count(src);
		Index **valid_array = (Index **) malloc(sizeof(Index *) * count);
		int i;
		for (i = 0; i < count; ++i) {
			valid_array[i] = (Index *) malloc(sizeof(Index));
		}

		get_index_block(src, valid_array, count);

		// sort the index array by res_offset.
		sort_index_block_by_res_offset(valid_array, count);
		// offset the res_offset.
		offset_index_block_res_offset(valid_array, count, sn, size_changed);
		// sort the index array by id.
		sort_index_block_by_id(valid_array, count);
		// write the index block to out file.
		write_index_block(valid_array, count, out);
		// free the space.
		for (i = 0; i < count; ++i) {
			free(valid_array[i]);
		}
		free(valid_array);
	} else {
		// get the size of rcc
		fseek(src, 0, SEEK_END);
		long src_size = ftell(src);

		// copy the rest bytes to out.
		if (copy_part(src, end_pos + 1, src_size - 1, out) == -1) {
			fseek(src, previous_pos, SEEK_SET);
			return -1;
		}
	}

	fseek(src, previous_pos, SEEK_SET);
	return 0;
}

int _replace_head(FILE *src, int sn, FILE *new_picture, FILE *out) {
	return 0;
}

int get_picture_size(FILE *src, unsigned int sn) {
	long current = ftell(src);

	long start = get_start_position(src, sn);
	// read the size
	int buffer;
	fseek(src, start - 4, SEEK_SET);
	fread((void*) &buffer, 4, 1, src);
	// the buffer got from file is little-endian, should be reversed.
	buffer = reverse_int32(buffer);
	fseek(src, current, SEEK_SET);
	return buffer;
}

void get_index_block(FILE *src, Index *index_array[], int count) {
	// seek to the start position
	int valid_length = get_valid_length(src);
	fseek(src, valid_length + 14, SEEK_SET);

	// read loop
	int i;
	for (i = 0; i < count; ++i) {
		index_array[i]->id = i;
		fread(&(index_array[i]->name_offset), 4, 1, src);
		index_array[i]->name_offset = reverse_int32(
				index_array[i]->name_offset);
		fseek(src, 6, SEEK_CUR);
		fread(&(index_array[i]->res_offset), 4, 1, src);
		index_array[i]->res_offset = reverse_int32(index_array[i]->res_offset);
	}
}

void sort_index_block_by_id(Index *index_array[], int count) {
	int i, j;
	Index *key;
	for (i = 1; i < count; ++i) {
		j = i - 1;
		key = index_array[i];
		while (j >= 0 && index_array[j]->id > key->id) {
			index_array[j + 1] = index_array[j];
			--j;
		}
		index_array[j + 1] = key;
	}
}

void sort_index_block_by_res_offset(Index *index_array[], int count) {
	int i, j;
	Index *key;
	for (i = 1; i < count; ++i) {
		j = i - 1;
		key = index_array[i];
		while (j >= 0 && index_array[j]->res_offset > key->res_offset) {
			index_array[j + 1] = index_array[j];
			--j;
		}
		index_array[j + 1] = key;
	}
}

void offset_index_block_res_offset(Index *index_array[], int count, int start,
		int offset) {
	int i;
	for (i = start; i < count; ++i) {
		index_array[i]->res_offset += offset;
	}
}

void write_index_block(Index *index_array[], int count, FILE *out) {
	int i;
	for (i = 0; i < count; ++i) {
		write_int(out, index_array[i]->name_offset, 0);
		write_int(out, 0, 0);
		write_short(out, 1, 0);
		write_int(out, index_array[i]->res_offset, 0);
	}
}
